package cn.iocoder.boot.springai001.provider;

import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.spec.McpSchema;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 动态MCP工具提供者
 * <p>
 * 基于官方文档的最佳实践，提供动态选择MCP客户端的功能
 * 参考: <a href="https://docs.spring.io/spring-ai/reference/api/mcp/mcp-helpers.html">...</a>
 * <p>
 * 使用场景：
 * 1. 根据用户权限动态加载不同的工具集
 * 2. 根据业务场景选择特定的MCP服务器
 * 3. 实现工具的按需加载和隔离
 */
@Slf4j
@Component
public class DynamicMcpToolProvider {

    private final List<McpSyncClient> mcpSyncClients;

    public DynamicMcpToolProvider(List<McpSyncClient> mcpSyncClients) {
        this.mcpSyncClients = mcpSyncClients;
        log.info("✅ 初始化动态MCP工具提供者，可用服务器数量: {}", mcpSyncClients.size());

        // 打印所有可用的服务器
        mcpSyncClients.forEach(client -> {
            McpSchema.Implementation info = client.getServerInfo();
            log.info("  - MCP服务器: {} (版本: {})", info.name(), info.version());
        });
    }

    /**
     * 根据服务器名称动态构建工具提供者
     * <p>
     * 官方文档示例：
     * <a href="https://docs.spring.io/spring-ai/reference/api/mcp/mcp-helpers.html">...</a>
     *
     * @param allowedServerNames 允许的服务器名称集合
     * @return 包含指定服务器工具的提供者
     */
    public ToolCallbackProvider buildProvider(Set<String> allowedServerNames) {
        log.info("构建动态工具提供者，允许的服务器: {}", allowedServerNames);

        // 按服务器名称过滤客户端
        List<McpSyncClient> selected = mcpSyncClients.stream()
                .filter(c -> allowedServerNames.contains(c.getServerInfo().name()))
                .toList();

        if (selected.isEmpty()) {
            log.warn("未找到匹配的MCP服务器: {}", allowedServerNames);
            return () -> new ToolCallback[0]; // 返回空提供者
        }

        log.info("成功筛选到 {} 个MCP服务器", selected.size());
        selected.forEach(client -> {
            String serverName = client.getServerInfo().name();
            try {
                int toolCount = client.listTools().tools().size();
                log.debug("  - {}: {} 个工具", serverName, toolCount);
            } catch (Exception e) {
                log.warn("获取服务器工具数量失败: {}", serverName);
            }
        });

        // 使用官方提供的 SyncMcpToolCallbackProvider
        return new SyncMcpToolCallbackProvider(selected);
    }

    /**
     * 获取所有工具提供者（不过滤）
     */
    public ToolCallbackProvider getAllToolsProvider() {
        log.info("构建完整工具提供者，包含所有 {} 个MCP服务器", mcpSyncClients.size());
        return new SyncMcpToolCallbackProvider(mcpSyncClients);
    }

    /**
     * 根据工具名称前缀过滤
     *
     * @param prefix 工具名称前缀（如 "filesystem_", "fetch_"）
     * @return 过滤后的工具回调列表
     */
    public List<ToolCallback> getToolsByPrefix(String prefix) {
        log.info("按前缀过滤工具: {}", prefix);

        List<ToolCallback> allCallbacks = SyncMcpToolCallbackProvider.syncToolCallbacks(mcpSyncClients);

        return allCallbacks.stream()
                .filter(callback -> callback.getToolDefinition().name().startsWith(prefix))
                .collect(Collectors.toList());
    }

    /**
     * 获取所有可用的服务器名称
     */
    public List<String> getAvailableServerNames() {
        return mcpSyncClients.stream()
                .map(c -> c.getServerInfo().name())
                .collect(Collectors.toList());
    }

    /**
     * 获取指定服务器的工具数量
     */
    public int getToolCount(String serverName) {
        return mcpSyncClients.stream()
                .filter(c -> c.getServerInfo().name().equals(serverName))
                .findFirst()
                .map(client -> {
                    try {
                        return client.listTools().tools().size();
                    } catch (Exception e) {
                        log.error("获取工具数量失败: {}", e.getMessage());
                        return 0;
                    }
                })
                .orElse(0);
    }

    /**
     * 示例：为不同用户角色提供不同的工具集
     *
     * @param userRole 用户角色（admin, developer, readonly）
     * @return 对应角色的工具提供者
     */
    public ToolCallbackProvider getProviderByRole(String userRole) {
        log.info("根据用户角色构建工具提供者: {}", userRole);

        Set<String> allowedServers = switch (userRole.toLowerCase()) {
            case "admin" -> Set.of("filesystem", "fetch", "excel", "everything");
            case "developer" -> Set.of("fetch", "excel");
            case "readonly" -> Set.of("fetch");
            default -> {
                log.warn("未知用户角色: {}，返回空工具集", userRole);
                yield Set.of();
            }
        };

        return buildProvider(allowedServers);
    }
}

